import torch

from ..builder import BBOX_ASSIGNERS
from ..iou_calculators import build_iou_calculator
from .assign_result import AssignResult
from .base_assigner import BaseAssigner
from typing import List
from .assigner_utils import print_bbox_in_img as _print_bbox_in_img
from .assigner_utils import print_num_anchor as _print_num_anchor
import cv2
import os

"""
wxz:
Atss_assigner_many_topk will allow you to use different topk in different fpn
level. For example, you can use topk=9 as the bottom level in fpn output to
select pos samples, and use topk=18 in top level.
"""


@BBOX_ASSIGNERS.register_module()
class ATSSAssigner_many_topk(BaseAssigner):
    """Assign a corresponding gt bbox or background to each bbox.

    Each proposals will be assigned with `0` or a positive integer
    indicating the ground truth index.

    - 0: negative sample, no assigned gt
    - positive integer: positive sample, index (1-based) of assigned gt

    Args:
        topk (float): number of bbox selected in each level
    """

    def __init__(
        self,
        topk,
        bottom_level_topk,
        iou_calculator=dict(type="BboxOverlaps2D"),
        ignore_iof_thr=-1,
        thre_method="mean+var",
        print_num_anchor=False,
        print_bbox_in_img=False,
        saved_path="./tmp_save/",
        add_lvl1_anchor_enable=False,
        add_lvl1_anchor_method="iou",
        add_lvl1_anchor_value=0.1,
        print_num_gt_without_anchor_center=False,
        print_pos_anchor_num=False,
    ):
        """init function of ATSS assigner.


        Args:
            topk: atss select topk candidate anchors in every FPN level except
                the bottom level.
            bottom_level_topk: atss threshold in the bottom FPN level.
            iou_calculator: which method being used to calculate the iou.
            ignore_iof_thr: ...
            extra_dict: many options designed by wxz in order to advance atss.
                thre_method: ["mean+var", "mean", "mean-var"]. which threshold
                    method used in atss, default: "mean+var".
                print_num_anchor: whether to print the number of positive anchor
                    assigned in every FPN level.
                print_bbox_in_img: whether to draw the positive anchor and
                    gt-bbox in image.
                saved_path: where to store the output of `print_num_anchor` and
                    `print_bbox_in_img`.
                add_lvl1_anchor: Due to the shrinking Stride=4 FPN featmap to
                    Stride=8 featmap, some small objects has less small positive
                    anchors in the lowest level(level 1) of fpn. So adding more
                    lvl1 positive anchors seems benefiting the converage of CNN
                    trainning.
                    - enable: whether add more lvl1 positive anchors.
                    - method: ["iou", "topk"]. which method to add more anchors.
                        "iou" means select those anchors whose iou are greater
                        than `value` below. Besides, if a gt has positive
                        anchors in other FPN featmap, it will be **disabled**.
                        "topk" means select top`thrshold` positive anchors for
                        every gt in the lvl1 FPN featmap.
                        Default is "iou".
                    - value:
                        This arg means iou value when "method" = "iou", and
                        "k" value when "method" = "topk".
        """
        self.topk = topk
        self.bottom_level_topk = bottom_level_topk
        self.iou_calculator = build_iou_calculator(iou_calculator)
        self.ignore_iof_thr = ignore_iof_thr

        self.thre_method = thre_method
        self.print_bbox_in_img = print_bbox_in_img
        self.print_num_anchor = print_num_anchor
        self.saved_path = saved_path

        self.add_lvl1_anchor_enable = add_lvl1_anchor_enable
        self.add_lvl1_anchor_method = add_lvl1_anchor_method
        self.add_lvl1_anchor_value = add_lvl1_anchor_value

        self.print_num_gt_without_anchor_center = print_num_gt_without_anchor_center
        self.num_gt_without_anchor_center = 0
        self.num_total_gt = 0

        self.print_pos_anchor_num = print_pos_anchor_num
        self.s4_pos_anchor_num = 0

        # https://github.com/sfzhang15/ATSS/blob/master/atss_core/modeling/rpn/atss/loss.py

    def assign(
        self,
        bboxes,
        num_level_bboxes,
        gt_bboxes,
        gt_bboxes_ignore=None,
        gt_labels=None,
        img_meta=None,
    ):
        """Assign gt to bboxes.

        The assignment is done in following steps

        1. compute iou between all bbox (bbox of all pyramid levels) and gt
        2. compute center distance between all bbox and gt
        3. on each pyramid level, for each gt, select k bbox whose center
           are closest to the gt center, so we total select k*l bbox as
           candidates for each gt
        4. get corresponding iou for the these candidates, and compute the
           mean and std, set mean + std as the iou threshold
        5. select these candidates whose iou are greater than or equal to
           the threshold as positive
        6. limit the positive sample's center in gt


        Args:
            bboxes (Tensor): Bounding boxes to be assigned, shape(n, 4).
            num_level_bboxes (List): num of bboxes in each level
            gt_bboxes (Tensor): Groundtruth boxes, shape (k, 4).
            gt_bboxes_ignore (Tensor, optional): Ground truth bboxes that are
                labelled as `ignored`, e.g., crowd boxes in COCO.
            gt_labels (Tensor, optional): Label of gt_bboxes, shape (k, ).

        Returns:
            :obj:`AssignResult`: The assign result.
        """

        INF = 100000000
        self.INF = INF
        # wxz: deepcopy bboxes
        bboxes = bboxes[:, :4]
        num_gt, num_bboxes = gt_bboxes.size(0), bboxes.size(0)

        # compute iou between all bbox and gt
        overlaps = self.iou_calculator(bboxes, gt_bboxes)

        # assign 0 by default
        assigned_gt_inds = overlaps.new_full((num_bboxes,), 0, dtype=torch.long)

        if num_gt == 0 or num_bboxes == 0:
            # No ground truth or boxes, return empty assignment
            max_overlaps = overlaps.new_zeros((num_bboxes,))
            if num_gt == 0:
                # No truth, assign everything to background
                assigned_gt_inds[:] = 0
            if gt_labels is None:
                assigned_labels = None
            else:
                assigned_labels = overlaps.new_full((num_bboxes,), -1, dtype=torch.long)
            return AssignResult(
                num_gt, assigned_gt_inds, max_overlaps, labels=assigned_labels
            )

        # compute center distance between all bbox and gt
        gt_cx = (gt_bboxes[:, 0] + gt_bboxes[:, 2]) / 2.0
        gt_cy = (gt_bboxes[:, 1] + gt_bboxes[:, 3]) / 2.0
        gt_points = torch.stack((gt_cx, gt_cy), dim=1)

        bboxes_cx = (bboxes[:, 0] + bboxes[:, 2]) / 2.0
        bboxes_cy = (bboxes[:, 1] + bboxes[:, 3]) / 2.0
        bboxes_points = torch.stack((bboxes_cx, bboxes_cy), dim=1)

        # if gt_points.shape = [2,2] ,gt_points[None,:,:].shape = [1,2,2]
        # this operation aims to expand dimension
        distances = (
            (bboxes_points[:, None, :] - gt_points[None, :, :]).pow(2).sum(-1).sqrt()
        )

        if (
            self.ignore_iof_thr > 0
            and gt_bboxes_ignore is not None
            # torch.numel->int. returns the total number of elements in the input tensor.
            and gt_bboxes_ignore.numel() > 0
            and bboxes.numel() > 0
        ):
            ignore_overlaps = self.iou_calculator(bboxes, gt_bboxes_ignore, mode="iof")
            ignore_max_overlaps, _ = ignore_overlaps.max(dim=1)
            ignore_idxs = ignore_max_overlaps > self.ignore_iof_thr
            distances[ignore_idxs, :] = INF
            assigned_gt_inds[ignore_idxs] = -1

        # Selecting candidates based on the center distance
        candidate_idxs = []
        start_idx = 0
        for level, bboxes_per_level in enumerate(num_level_bboxes):
            # on each pyramid level, for each gt,
            # select k bbox whose center are closest to the gt center
            end_idx = start_idx + bboxes_per_level
            distances_per_level = distances[start_idx:end_idx, :]
            if level == 0:
                selectable_k = min(self.bottom_level_topk, bboxes_per_level)
            else:
                selectable_k = min(self.topk, bboxes_per_level)
            # wxz: see https://pytorch.org/docs/stable/generated/torch.topk.html#torch.topk
            _, topk_idxs_per_level = distances_per_level.topk(
                selectable_k, dim=0, largest=False
            )
            candidate_idxs.append(topk_idxs_per_level + start_idx)
            start_idx = end_idx
        candidate_idxs = torch.cat(candidate_idxs, dim=0)
        candidate_overlaps = overlaps[candidate_idxs, torch.arange(num_gt)]
        overlaps_mean_per_gt = candidate_overlaps.mean(0)
        overlaps_std_per_gt = candidate_overlaps.std(0)
        if self.thre_method == "mean+var":
            overlaps_thr_per_gt = overlaps_mean_per_gt + overlaps_std_per_gt
        elif self.thre_method == "mean":
            overlaps_thr_per_gt = overlaps_mean_per_gt
        elif self.thre_method == "mean-var":
            overlaps_thr_per_gt = overlaps_mean_per_gt - overlaps_std_per_gt
        else:
            raise ValueError(f"No such '{self.thre_method}' thre_method.")

        is_pos = candidate_overlaps >= overlaps_thr_per_gt[None, :]

        if self.add_lvl1_anchor_enable:
            # add some positive anchors in the lowest level of FPN.

            if self.add_lvl1_anchor_method == "iou":
                lvl1_thr = torch.full_like(
                    overlaps_thr_per_gt, self.add_lvl1_anchor_value
                )
                lvl1_overlaps = candidate_overlaps.clone()
                lvl1_overlaps[selectable_k:] = 0

                # 只允许在第一层以外其他层没有正样本的gt能够得到更多正样本
                has_pos_in_other_layer = is_pos[selectable_k:].sum(dim=0, dtype=bool)
                lvl1_overlaps = lvl1_overlaps * ~has_pos_in_other_layer

                lvl1_pos = lvl1_overlaps >= lvl1_thr[None, :]

            elif self.add_lvl1_anchor_method == "topk":
                lvl1_overlaps = candidate_overlaps.clone()
                lvl1_overlaps[selectable_k:] = 0
                _, indice = lvl1_overlaps.topk(self.add_lvl1_anchor_value, dim=0)

                # 需要传入两个一维tensor做索引。分别构建横纵坐标。
                indice = indice.permute(1, 0).reshape(-1)
                extra_indice = [i for i in range(num_gt)] * self.add_lvl1_anchor_value
                extra_indice = torch.tensor(extra_indice)

                lvl1_pos = torch.full_like(lvl1_overlaps, 0, dtype=bool)
                lvl1_pos[indice, extra_indice] = True

            else:
                raise ValueError(
                    f"no such add_lvl1_anchor method: {self.add_lvl1_anchor_method}"
                )

            is_pos = is_pos | lvl1_pos

        # limit the positive sample's center in gt
        # TODO(wangxianzhuo) why increase the index of candidate? ---------
        for gt_idx in range(num_gt):
            candidate_idxs[:, gt_idx] += gt_idx * num_bboxes
        # -----------------------------------------------------------------
        ep_bboxes_cx = (
            bboxes_cx.view(1, -1).expand(num_gt, num_bboxes).contiguous().view(-1)
        )
        ep_bboxes_cy = (
            bboxes_cy.view(1, -1).expand(num_gt, num_bboxes).contiguous().view(-1)
        )
        candidate_idxs = candidate_idxs.view(-1)

        # calculate the left, top, right, bottom distance between positive
        # bbox center and gt side
        l_ = ep_bboxes_cx[candidate_idxs].view(-1, num_gt) - gt_bboxes[:, 0]
        t_ = ep_bboxes_cy[candidate_idxs].view(-1, num_gt) - gt_bboxes[:, 1]
        r_ = gt_bboxes[:, 2] - ep_bboxes_cx[candidate_idxs].view(-1, num_gt)
        b_ = gt_bboxes[:, 3] - ep_bboxes_cy[candidate_idxs].view(-1, num_gt)
        is_in_gts = torch.stack([l_, t_, r_, b_], dim=1).min(dim=1)[0] > 0.01
        is_pos = is_pos & is_in_gts

        # see how many gts do not contain even a single anchor center.
        if self.print_num_gt_without_anchor_center:
            print("--- the number of gt which contains no anchor-centers---")
            has_anchor_center_in_gt = is_in_gts[:].sum(dim=0, dtype=bool)

            self.num_gt_without_anchor_center += (
                has_anchor_center_in_gt == False
            ).sum()
            self.num_total_gt += num_gt
            print(f"{img_meta['filename']}")
            print(
                f"no anchor gt / all gt :{self.num_gt_without_anchor_center} / {self.num_total_gt}"
            )

        # if an anchor box is assigned to multiple gts,
        # the one with the highest IoU will be selected.
        # wxz: https://pytorch.org/docs/stable/generated/torch.t.html#torch.t
        # wxz: t: transpose
        overlaps_inf = torch.full_like(overlaps, -INF).t().contiguous().view(-1)
        index = candidate_idxs.view(-1)[is_pos.view(-1)]
        overlaps_inf[index] = overlaps.t().contiguous().view(-1)[index]
        overlaps_inf = overlaps_inf.view(num_gt, -1).t()

        max_overlaps, argmax_overlaps = overlaps_inf.max(dim=1)
        assigned_gt_inds[max_overlaps != -INF] = (
            argmax_overlaps[max_overlaps != -INF] + 1
        )

        #  saved_path = "/workspace/volume/wangxianzhuo-data/data/facecar_visual/anchor/0"
        if self.print_bbox_in_img:
            os.makedirs(self.saved_path, exist_ok=True)
            _print_bbox_in_img(
                img_meta, gt_bboxes, bboxes, max_overlaps, self.saved_path
            )

        if self.print_num_anchor:
            os.makedirs(self.saved_path, exist_ok=True)
            _print_num_anchor(
                bboxes,
                gt_bboxes,
                num_bboxes,
                num_level_bboxes,
                index,
                img_meta,
                gt_labels,
                self.saved_path,
                overlaps_thr_per_gt,
            )

        if gt_labels is not None:
            assigned_labels = assigned_gt_inds.new_full((num_bboxes,), -1)
            pos_inds = torch.nonzero(assigned_gt_inds > 0, as_tuple=False).squeeze()
            if pos_inds.numel() > 0:
                assigned_labels[pos_inds] = gt_labels[assigned_gt_inds[pos_inds] - 1]

                # print the number of pos anchors in the bottom level
                if self.print_pos_anchor_num:
                    self.s4_pos_anchor_num += (
                        (pos_inds <= num_level_bboxes[0]).sum().tolist()
                    )
                    print(
                        f"pos anchor amount in the bottom level : {self.s4_pos_anchor_num}"
                    )
        else:
            assigned_labels = None

        return AssignResult(
            num_gt, assigned_gt_inds, max_overlaps, labels=assigned_labels
        )
